import os

import cv2


def load_data(path):
    img_list = []
    label_list = []
    dir_counter = 2
    for img_path in os.listdir(path):
        img_name = img_path
        img_label = img_path.split(".", 1)
        img_path = path + "/" + img_path
        Suffix_name = ['.png', '.jpg', '.jpeg']
        if img_path.endswith(tuple(Suffix_name)):
            img = cv2.imread(img_path)
            # img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            img = cv2.resize(img,(150,150))
            # cv2.imshow("",img)
            # print(img_label[0])
            # cv2.waitKey(100)
            img_list.append(img)
            if img_label[0] == 'cat':
                cv2.imwrite(path+r'/cat/'+img_name,img)
                # label_list.append(0)
            else:
                cv2.imwrite(path + r'/dog/' + img_name,img)
                # label_list.append(1)
    return img_list, label_list, dir_counter

img_list, label_list, dir_counter = load_data(path="train")
print(label_list)
print(dir_counter)